{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.4.2 召回内容上下文扩充"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.storage import InMemoryStore\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.retrievers.parent_document_retriever import ParentDocumentRetriever\n",
    "loaders = [\n",
    "    TextLoader(\"./file1.txt\"),\n",
    "    TextLoader(\"./file2.txt\"),\n",
    "]\n",
    "docs = []\n",
    "for l in loaders:\n",
    "    docs.extend(l.load())\n",
    "\n",
    "parent_splitter = RecursiveCharacterTextSplitter(chunk_size=1000)\n",
    "child_splitter = RecursiveCharacterTextSplitter(chunk_size=100)\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"split_parents\", embedding_function=OpenAIEmbeddings()\n",
    ")\n",
    "store = InMemoryStore()\n",
    "retriever = ParentDocumentRetriever(\n",
    "    child_splitter=child_splitter,\n",
    "    parent_splitter=parent_splitter,\n",
    "    vectorstore=vectorstore,\n",
    "    docstore=store,\n",
    ")\n",
    "retriever.add_documents(docs)\n",
    "retrieved_docs = retriever.get_relevant_documents(\"西瓜的品种\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.4.3 文本多向量表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.storage import InMemoryStore\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "import uuid\n",
    "import os\n",
    "os.environ[\"OPENAI_API_BASE\"] = \"xxx\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"xxx\"\n",
    "\n",
    "loader = TextLoader(\"./file.txt\")\n",
    "docs = []\n",
    "docs.extend(loader.load())\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000,chunk_overlap=0)\n",
    "docs = text_splitter.split_documents(docs)\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"full_documents\", embedding_function=OpenAIEmbeddings()\n",
    ")\n",
    "\n",
    "store = InMemoryStore()\n",
    "# metadata中，储存大块文本id的键\n",
    "id_key = \"doc_id\"\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    docstore=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "# 生成大块文本对应的id值\n",
    "doc_ids = [str(uuid.uuid4()) for _ in docs]\n",
    "child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000,chunk_overlap=0)\n",
    "sub_docs = []\n",
    "for i, doc in enumerate(docs):\n",
    "    _id = doc_ids[i]\n",
    "    # 大块文本进一步切分成小块文本\n",
    "    _sub_docs = child_text_splitter.split_documents([doc])\n",
    "    for _doc in _sub_docs:\n",
    "        # 小块文本的metadata中保存其对应的大块文本的id\n",
    "        _doc.metadata[id_key] = _id\n",
    "    sub_docs.extend(_sub_docs)\n",
    "# 向量数据库中存储的是小块文本\n",
    "retriever.vectorstore.add_documents(sub_docs)\n",
    "# 存储大块文本id和内容的对应关系\n",
    "retriever.docstore.mset(list(zip(doc_ids, docs)))\n",
    "# 召回小块文本\n",
    "print(len(retriever.vectorstore.similarity_search(\"user query\")[0].page_content))\n",
    "# 召回小块文本对应的大块文本\n",
    "print(len(retriever.get_relevant_documents(\"user query\")[0].page_content))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.4.4 查询内容优化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a sample vectorDB\n",
    "from langchain.document_loaders import WebBaseLoader\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
    "import logging\n",
    "\n",
    "# Load blog post\n",
    "loader = WebBaseLoader(\"https://lilianweng.github.io/posts/2023-06-23-agent/\")\n",
    "data = loader.load()\n",
    "\n",
    "# Split\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)\n",
    "splits = text_splitter.split_documents(data)\n",
    "\n",
    "# VectorDB\n",
    "embedding = OpenAIEmbeddings()\n",
    "vectordb = Chroma.from_documents(documents=splits, embedding=embedding)\n",
    "\n",
    "question = \"What are the approaches to Task Decomposition?\"\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "retriever_from_llm = MultiQueryRetriever.from_llm(\n",
    "    retriever=vectordb.as_retriever(), llm=llm\n",
    ")\n",
    "# Set logging for the queries\n",
    "logging.basicConfig()\n",
    "logging.getLogger(\"langchain.retrievers.multi_query\").setLevel(logging.INFO)\n",
    "\n",
    "unique_docs = retriever_from_llm.get_relevant_documents(query=question)\n",
    "len(unique_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import HypotheticalDocumentEmbedder \n",
    "from langchain.embeddings import OpenAIEmbeddings \n",
    "from langchain.llms import OpenAI \n",
    "import os \n",
    "base_embeddings = OpenAIEmbeddings() \n",
    "llm = OpenAI() \n",
    "embeddings = HypotheticalDocumentEmbedder.from_llm(llm, base_embeddings, \"web_search\") \n",
    "result = embeddings.embed_query(\"北京在哪儿\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.4.5 召回文本重排序"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import LLMChainExtractor,LLMChainFilter\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_BASE\"]  = \"xxx\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"xxx\"\n",
    "documents = TextLoader('./fruit_information.txt').load()\n",
    "text_splitter=CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "retriever=FAISS.from_documents(texts, OpenAIEmbeddings()).as_retriever()\n",
    "llm = OpenAI(temperature=0)\n",
    "# 提取关键句的压缩器\n",
    "# compressor = LLMChainExtractor.from_llm(llm)\n",
    "# 选择召回文本的压缩器\n",
    "compressor = LLMChainFilter.from_llm(llm)\n",
    "compression_retriever=ContextualCompressionRetriever(\n",
    "base_compressor=compressor, base_retriever=retriever)\n",
    "compressed_docs = compression_retriever.get_relevant_documents(\n",
    "\"介绍一下葡萄\")\n",
    "for d in compressed_docs:\n",
    "    print(d.page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.4.6 多检索器融合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import BM25Retriever, EnsembleRetriever\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "import os\n",
    "\n",
    "doc_list = [\n",
    "    \"我喜欢吃西瓜\",\n",
    "    \"我喜欢吃葡萄\",\n",
    "    \"葡萄和西瓜都是我喜欢的水果\",\n",
    "]\n",
    "\n",
    "# 定义BM25检索器\n",
    "bm25 = BM25Retriever.from_texts(doc_list)\n",
    "bm25.k = 2\n",
    "\n",
    "#定义基于OpenAI向量化模型的检索器\n",
    "embedding = OpenAIEmbeddings(\n",
    "    openai_api_base=\"xxx\",\n",
    "    openai_api_key=\"xxx\"\n",
    ")\n",
    "faiss_vectorstore = FAISS.from_texts(doc_list, embedding)\n",
    "embedding_retriever = faiss_vectorstore.as_retriever(search_kwargs={\"k\": 2})\n",
    "\n",
    "# 检索器融合\n",
    "ensemble_retriever = EnsembleRetriever(\n",
    "    retrievers=[bm25, embedding_retriever], weights=[0.2, 0.8]\n",
    ")\n",
    "docs = ensemble_retriever.get_relevant_documents(\"葡萄\")\n",
    "print(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4.4.7 结合元数据召回"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import FAISS\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "import os\n",
    "from langchain.docstore.document import Document\n",
    "doc_list = [\n",
    "    Document(page_content=\"葡萄价格：10元/kg\",metadata={\"shop_name\":\"水果店1\"}),\n",
    "    Document(page_content=\"西瓜价格：5元/kg\",metadata={\"shop_name\":\"水果店1\"}),\n",
    "    Document(page_content=\"葡萄价格：15元/kg\",metadata={\"shop_name\":\"水果店2\"}),  \n",
    "]\n",
    "#定义基于OpenAI向量化模型的检索器\n",
    "embedding = OpenAIEmbeddings(\n",
    "    openai_api_base=\"xxx\",\n",
    "    openai_api_key=\"xxx\"\n",
    ")\n",
    "# 带元数据过滤的检索\n",
    "faiss_vectorstore_filter = FAISS.from_documents(doc_list, embedding)\n",
    "faiss_retriever_filter = faiss_vectorstore_filter.as_retriever(\n",
    "    search_kwargs={\"k\": 2, \"filter\":{\"shop_name\":\"水果店1\"}})\n",
    "print(faiss_retriever_filter.get_relevant_documents(\"葡萄价格\"))\n",
    "# 输出===\n",
    "#[Document(page_content='葡萄价格：10元/kg', metadata={'shop_name': '水果店1'}), \n",
    "#Document(page_content='西瓜价格：5元/kg', metadata={'shop_name': '水果店1'})]\n",
    "#===\n",
    "\n",
    "# 不带元数据过滤的检索\n",
    "faiss_vectorstore = FAISS.from_documents(doc_list, embedding)\n",
    "faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={\"k\": 2})\n",
    "print(faiss_retriever.get_relevant_documents(\"葡萄价格\"))\n",
    "# 输出===\n",
    "#[Document(page_content='葡萄价格：15元/kg', metadata={'shop_name': '水果店2'}), \n",
    "#Document(page_content='葡萄价格：10元/kg', metadata={'shop_name': '水果店1'})]\n",
    "#==="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.document_transformers.openai_functions import create_metadata_tagger\n",
    "from langchain.schema import Document\n",
    "import json\n",
    "schema = {\n",
    "    \"properties\": {\n",
    "        \"fruit\": {\"type\": \"string\"},\n",
    "        \"shop_name\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": \"the name of fruit shop\",\n",
    "        },\n",
    "    },\n",
    "    \"required\": [\"fruit\", \"shop_name\"],\n",
    "}\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(openai_api_base=\"xxx\",\n",
    "    openai_api_key=\"xxx\",temperature=0)\n",
    "\n",
    "gen_tagger = create_metadata_tagger(metadata_schema=schema, llm=llm)\n",
    "documents = [\n",
    "    Document(\n",
    "        page_content=\"味道美水果店卖葡萄，3块钱一斤\"\n",
    "    ),\n",
    "]\n",
    "\n",
    "enhanced_documents = gen_tagger.transform_documents(documents)\n",
    "\n",
    "for d in enhanced_documents:\n",
    "    print(d.metadata)\n",
    "#输出===\n",
    "#{'fruit': '葡萄', 'shop_name': '味道美水果店'}\n",
    "#==="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from datetime import datetime, timedelta\n",
    "from langchain.docstore import InMemoryDocstorefrom \n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.retrievers import TimeWeightedVectorStoreRetriever\n",
    "from langchain.schema import Document\n",
    "import os\n",
    "os.environ[\"OPENAI_API_BASE\"]  = \"xxx\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"xxx\"\n",
    "# 定义向量化模型\n",
    "embeddings_model = OpenAIEmbeddings()\n",
    "# 构建faiss向量数据库\n",
    "embedding_size = 1536\n",
    "index = faiss.IndexFlatL2(embedding_size)\n",
    "vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})\n",
    "# 定义时间衰减检索器\n",
    "retriever = TimeWeightedVectorStoreRetriever(vectorstore=vectorstore, decay_rate=0.999, k=1)\n",
    "yesterday = datetime.now() - timedelta(days=1)\n",
    "retriever.add_documents([Document(page_content=\"葡萄价格3元/斤\", metadata={\"last_accessed_at\": yesterday})])\n",
    "retriever.add_documents([Document(page_content=\"西瓜价格2元/斤\")])\n",
    "# \"Hello Foo\" is returned first because \"hello world\" is mostly forgotten\n",
    "docs = retriever.get_relevant_documents(\"查询葡萄价格\")\n",
    "print(docs)\n",
    "#输出===\n",
    "#[Document(page_content='西瓜价格2元/斤', metadata={'last_accessed_at': datetime.datetime(2023, 6, 18, 15, 19, 54, 840124), 'created_at': datetime.datetime(2023, 6, 18, 815, 19, 53, 640318), 'buffer_idx': 1})]\n",
    "#==="
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
